#!/usr/bin/env python

"""
直接键盘控制方案 - 通过控制台输入实现

这个脚本提供了一个更直接的方法来解决键盘输入无响应的问题：
1. 使用 pynput 库从控制台监听键盘事件
2. 在主循环中处理键盘输入
3. 手动更新机器人位置

使用方法：
python direct_keyboard_control.py --config_path lerobot/configs/gym_hil_env.json
"""

import os
import sys
import time
import json
import argparse
import numpy as np
import torch
from pathlib import Path
from pynput import keyboard

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent.parent
sys.path.append(str(project_root))

# 导入必要的模块
from lerobot.common.envs.configs import EnvConfig
from lerobot.configs.parser import wrap
from lerobot.scripts.rl.gym_manipulator import make_robot_env

# 全局变量存储按键状态
key_states = {
    'up': False,
    'down': False,
    'left': False,
    'right': False,
    'shift_l': False,
    'shift_r': False,
    'ctrl_l': False,
    'ctrl_r': False,
    'enter': False,
    'backspace': False,
    'space': False,
    'esc': False
}

# 全局变量控制程序运行
running = True

def parse_args():
    parser = argparse.ArgumentParser(description="直接键盘控制机器人")
    parser.add_argument("--config_path", type=str, required=True, help="配置文件路径")
    return parser.parse_args()

# 按键按下时的回调函数
def on_press(key):
    global key_states, running
    try:
        if key == keyboard.Key.up:
            key_states['up'] = True
        elif key == keyboard.Key.down:
            key_states['down'] = True
        elif key == keyboard.Key.left:
            key_states['left'] = True
        elif key == keyboard.Key.right:
            key_states['right'] = True
        elif key == keyboard.Key.shift:
            key_states['shift_l'] = True
        elif key == keyboard.Key.shift_r:
            key_states['shift_r'] = True
        elif key == keyboard.Key.ctrl_l:
            key_states['ctrl_l'] = True
        elif key == keyboard.Key.ctrl_r:
            key_states['ctrl_r'] = True
        elif key == keyboard.Key.enter:
            key_states['enter'] = True
        elif key == keyboard.Key.backspace:
            key_states['backspace'] = True
        elif key == keyboard.Key.space:
            key_states['space'] = True
        elif key == keyboard.Key.esc:
            key_states['esc'] = True
            running = False  # 设置退出标志
    except AttributeError:
        pass

# 按键释放时的回调函数
def on_release(key):
    global key_states
    try:
        if key == keyboard.Key.up:
            key_states['up'] = False
        elif key == keyboard.Key.down:
            key_states['down'] = False
        elif key == keyboard.Key.left:
            key_states['left'] = False
        elif key == keyboard.Key.right:
            key_states['right'] = False
        elif key == keyboard.Key.shift:
            key_states['shift_l'] = False
        elif key == keyboard.Key.shift_r:
            key_states['shift_r'] = False
        elif key == keyboard.Key.ctrl_l:
            key_states['ctrl_l'] = False
        elif key == keyboard.Key.ctrl_r:
            key_states['ctrl_r'] = False
        elif key == keyboard.Key.enter:
            key_states['enter'] = False
        elif key == keyboard.Key.backspace:
            key_states['backspace'] = False
        elif key == keyboard.Key.space:
            key_states['space'] = False
        elif key == keyboard.Key.esc:
            key_states['esc'] = False
    except AttributeError:
        pass

@wrap()
def main(cfg: EnvConfig):
    global running, key_states
    
    # 创建环境
    env = make_robot_env(cfg)
    
    # 重置环境
    obs, info = env.reset()
    
    # 定义步长
    # 从配置中提取步长参数
    if hasattr(cfg, 'wrapper') and hasattr(cfg.wrapper, 'ee_action_space_params'):
        step_size = {
            "x": cfg.wrapper.ee_action_space_params.x_step_size,
            "y": cfg.wrapper.ee_action_space_params.y_step_size,
            "z": cfg.wrapper.ee_action_space_params.z_step_size
        }
    else:
        # 默认步长
        step_size = {"x": 0.01, "y": 0.01, "z": 0.01}
    
    # 打印控制说明
    print("键盘控制说明:")
    print("  方向键: 在 X-Y 平面移动")
    print("  Shift 和 右Shift: 在 Z 轴移动")
    print("  右Ctrl 和 左Ctrl: 打开和关闭夹爪")
    print("  Enter: 成功结束回合")
    print("  Backspace: 失败结束回合")
    print("  空格: 开始/停止干预")
    print("  ESC: 退出")
    print("\n注意: 控制台必须保持焦点以捕获键盘输入\n")
    
    # 启动键盘监听器（非阻塞）
    listener = keyboard.Listener(on_press=on_press, on_release=on_release)
    listener.start()
    
    print("键盘监听已启动，可以开始控制机器人...")
    
    # 主循环
    while running:
        # 打印当前按键状态（调试用）
        active_keys = [k for k, v in key_states.items() if v]
        if active_keys:
            print(f"当前按下的键: {active_keys}")
        
        # 计算动作
        action = np.zeros(7)  # [dx, dy, dz, 夹爪]
        
        # X-Y 平面移动
        if key_states['up']:
            action[0] += step_size["x"]
        if key_states['down']:
            action[0] -= step_size["x"]
        if key_states['left']:
            action[1] += step_size["y"]
        if key_states['right']:
            action[1] -= step_size["y"]
        
        # Z 轴移动
        if key_states['shift_l']:
            action[2] -= step_size["z"]
        if key_states['shift_r']:
            action[2] += step_size["z"]
        
        # 夹爪控制
        if key_states['ctrl_l']:
            action[3] = -1  # 关闭夹爪
        elif key_states['ctrl_r']:
            action[3] = 1   # 打开夹爪
        
        # 检查是否有按键被按下，只有在有按键时才执行动作
        if any(key_states.values()):
            # 将 NumPy 数组转换为 PyTorch 张量
            # 默认使用 CPU 设备，除非配置中指定了其他设备
            device = getattr(cfg, 'device', 'cpu')
            action_tensor = torch.tensor(action, dtype=torch.float32, device=device)
            
            # 执行动作
            obs, reward, terminated, truncated, info = env.step(action_tensor)
            
            # 如果回合结束，重置环境
            if terminated or truncated or key_states['enter'] or key_states['backspace']:
                print("回合结束，重置环境...")
                obs, info = env.reset()
        
        # 限制帧率
        time.sleep(0.1)
    
    # 清理资源
    print("正在关闭环境...")
    env.close()
    listener.stop()
    print("程序已退出")

if __name__ == "__main__":
    main()
